# 导包
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras import Model
from tensorflow.keras.layers import Flatten, Conv2D, MaxPool2D, Dense

"""
卷积操作tf.keras.layers.Conv2D
用到的参数:
- input_shape:输入数据的维度大小，如3通道64X64的输入图片，要设置为（64，64，3）
参数input_shape只需要在第一个网络层进行设置。
- filters:输入整数，卷积核个数（等于卷积后输出的通道数）
- kernel_size:卷积核的大小，通常为一个正整数或者一个列表。
  例子: kernel_szie = 3 ,kernel_size(3,3),都表示卷积核宽高方向的大小为3
- strides:卷积核的滑动步长，通常为一个正整数或者一个列表，列表可以包含2个或者4个元素，默认为（1，1）
  例子: strides=2,strides=(2,2),都表示宽高方向的滑动步长为2
- padding:补零设置，可选项为”VALID“,"SAME"。
 - if padding="VALID",表示不做补零操作。
 - if padding="SAME",程序会自动补零，使得输出feature map的宽高=ceil(输入宽高/stride)
  ceil表示向上取整，当strides=1时，输入输出的宽高是一致的。
- activation:激活函数
  可选项:
  -"sigmoid":sigmoid激活函数
  -”tanh":tanh激活函数
  -“relu":relu激活函数
  -”elu"或者tf.keras.activations.elu（alpha=1.0):elu激活函数
  -“selu":selu激活函数
  -”swish":swish激活函数（tf2.2版本以上才有）
  -“softmax":softmax函数
  -input_shape:如果是第一层卷积，需要设置输入图片的大小（height,width,channel)
  比如input_shape(128,128,3)
  -name 输入字符串，给该层设置一个名称
  
池化操作tf.keras.layer.Maxpool2D
用到的参数：
- pool_size:池化卷积核的大小，具体设置与tf.keras.layers.Conv2D的kernel_size一样
- strides:池化的滑动步长
- name:输入字符串，给该层设置一个名称

全连接操作tf.keras.layers.Dense
用到的参数
- units:输入整数，全连接层神经元个数，
- activation:激活函数，分类网络的输出层一般用'softmax'激活函数
- name

- 展平操作 tf.keras.layers.Flateen

"""

# 定义一个子类来搭建模型
"""
class ConvModel(Model):
    def __init__(self):
        # 父类初始化
        super(ConvModel, self).__init__()
        # 卷积层 conv_1_1
        self.conv_1_1 = Conv2D(input_shape=(64, 64, 3), filters=32, kernel_size=3, padding='same', activation='relu',
                               name='conv_1_1')
        # 卷积层 conv_1_2
        self.conv_1_2 = Conv2D(filters=32, kernel_size=3, padding='same', activation='relu', name='conv_1_2')
        # 池化层 max_pool_1
        self.max_pool_1 = MaxPool2D(pool_size=2, name='max_pool_1')
        # 卷积层 conv_2_1 = Conv2D()
        self.conv_2_1 = Conv2D(filters=64, kernel_size=3, padding='same', activation='relu', name='conv_2_1')
        # 卷积层 conv_2.2
        self.conv_2_2 = Conv2D(filters=64, kernel_size=3, padding='same', activation='relu', name='conv_2_2')
        # 池化层 max_pool_2
        self.max_pool_2 = MaxPool2D(pool_size=2, name='max_pool_2')
        # 展平层flatten
        self.flatten = Flatten(name='flatten')
        # 全连接层
        self.dense = Dense(units=6, activation='softmax', name='logit')
"""


class ConvModel(Model):
    def __init__(self):
        # 父类初始化
        super(ConvModel, self).__init__()
        # 卷积层 conv_1_1
        self.conv_1_1 = Conv2D(input_shape=(160, 160, 3), filters=128, kernel_size=11, strides=4, padding='same',
                               activation='relu', name='conv_1_1')
        # 卷积层 conv_1_2
        self.conv_1_2 = Conv2D(filters=256, kernel_size=5, padding='same', activation='relu', name='conv_1_2')
        # 池化层 max_pool_1
        self.max_pool_1 = MaxPool2D(pool_size=3, strides=2, name='max_pool_1')
        # 卷积层 conv_2_1 = Conv2D()
        self.conv_2_1 = Conv2D(filters=384, kernel_size=3, padding='same', activation='relu', name='conv_2_1')
        # 卷积层 conv_2.2
        self.conv_2_2 = Conv2D(filters=384, kernel_size=3, padding='same', activation='relu', name='conv_2_2')
        # 卷积层 conv_2.3
        self.conv_2_3 = Conv2D(filters=256, kernel_size=3, padding='same', activation='relu', name='conv_2_2')
        # 池化层 max_pool_2
        self.max_pool_2 = MaxPool2D(pool_size=3, strides=2, name='max_pool_2')
        # 展平层flatten
        self.flatten = Flatten(name='flatten')
        # 全连接层
        self.dense = Dense(units=10, activation='softmax', name='logit')

    def call(self, x):
        x = self.conv_1_1(x)
        x = self.conv_1_2(x)
        x = self.max_pool_1(x)
        x = self.conv_2_1(x)
        x = self.conv_2_2(x)
        x = self.conv_2_3(x)
        x = self.max_pool_2(x)
        x = self.flatten(x)
        x = self.dense(x)
        return x
